from sql_options import get_data_clients_data
import os
import tenseal as ts
from pypbc import *
import hashlib
Hash = hashlib.sha256
from utils.masked_input import Flatten_paremeter, KeyGen, PPGen, Hhash, sign_key_gen, sign, Verify_agg_sign_2, \
    decompose_agg_w_new, \
    Hhash_equa, Verify_agg_sign_1, Restore_values_to_w, Invert_Quality

import base64

server = None
class ServerAgg:
    def __init__(self, args,data):
        self.args = args.copy()
        self.device = args["device"]
        self.context = None
        self.params = Parameters(param_string=data['params'])
        self.pairing = Pairing(self.params)
        self.G2 = G2
        self.g2 = Element(self.pairing, G2, value=data['g2'])
        self.total_params = data['total_params']
        self.num_clients = data['num_clients']
        context_bytes = base64.b64decode(data['context'])
        self.context = ts.context_from(context_bytes)
        self.pk_sign_2_s = [0 for i in range(self.num_clients)]

    def aggregation(self,server_state,connection):
        enc_w_new_s_recover, pk_sign_1_s, sig1_s_recover, sig2_s_recover = [], [], [], []
        local_BFV_bytes_base64 = None
        if self.args["method"] == 'VeriTrac':
            all_clients = get_data_clients_data(connection)
            for client in all_clients:
                client_num_id = client[0]
                client_sign1 = client[1]
                client_sign2 = client[2]
                client_BFV = client[3]
                pk_sign_1 = client[4]
                pk_sign_2 = client[5]
                data_source = os.getenv('DB_DATA_SOURCE', 'mysql')
                if data_source != 'mysql':
                    client_BFV = client_BFV.tobytes()
                enc_w_new_s_recover.append(client_BFV)
                print("client_BFV length", len(client_BFV))
                sig1_s_recover.append(client_sign1)
                sig2_idx_recover = Element(self.pairing, self.G2, value=client_sign2)
                sig2_s_recover.append(sig2_idx_recover)
                pk_sign_1_s.append(pk_sign_1)
                pk_sign_2 = Element(self.pairing, self.G2, value=pk_sign_2)
                self.pk_sign_2_s[client_num_id-1] = pk_sign_2

            verify_aggregation = Verify_agg_sign_2(enc_w_new_s_recover, self.pairing, sig1_s_recover, self.pk_sign_2_s,
                                                   self.num_clients, self.g2, sig2_s_recover)
            print("verify_aggregation = ", verify_aggregation)

            if verify_aggregation == 1 and server_state == True:

                len_h = 130 * self.num_clients
                h_ini = [0] * len_h

                w_local_forged = h_ini +[1]*self.total_params
                enc_w_local_forged = ts.bfv_vector(self.context, w_local_forged)
                w_new_s_BFV = ts.bfv_vector_from(self.context, enc_w_new_s_recover[self.num_clients - 1])  # 反序列化
                w_new_s_BFV1= w_new_s_BFV
                for i in range(self.num_clients - 1):
                    w_new_s_BFV = w_new_s_BFV + ts.bfv_vector_from(self.context, enc_w_new_s_recover[i])
                w_new_s_BFV = w_new_s_BFV + enc_w_local_forged

                local_BFV_bytes = w_new_s_BFV.serialize()
                local_BFV_bytes_base64 = base64.b64encode(local_BFV_bytes).decode('utf-8')
            elif verify_aggregation == 1 and server_state == False:
                w_new_s_BFV = ts.bfv_vector_from(self.context, enc_w_new_s_recover[self.num_clients - 1])  # 反序列化
                for i in range(self.num_clients - 1):
                    w_new_s_BFV = w_new_s_BFV + ts.bfv_vector_from(self.context, enc_w_new_s_recover[i])
                local_BFV_bytes = w_new_s_BFV.serialize()
                local_BFV_bytes_base64=base64.b64encode(local_BFV_bytes).decode('utf-8')
                print("local_BFV_bytes_base64 length", len(local_BFV_bytes_base64))

        return verify_aggregation , local_BFV_bytes_base64, pk_sign_1_s, sig1_s_recover
